package tech.waterism.topo.calculator;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import tech.waterism.topo.context.CalculationContext;
import tech.waterism.topo.model.*;

import java.util.*;

import static org.junit.jupiter.api.Assertions.*;

/**
 * ForecastCalculator 单元测试
 * 
 * @author waterism
 * @since 1.0.0
 */
class ForecastCalculatorTest {
    
    private ForecastCalculator calculator;
    private CalculationContext context;
    
    @BeforeEach
    void setUp() {
        calculator = new ForecastCalculator();
        context = new CalculationContext();
        context.setStartTime(new Date(System.currentTimeMillis() - 3600000)); // 1小时前
    }
    
    @Test
    void testCalculate_EmptyInputs() {
        // 准备：没有输入分量的节点
        TopoNode node = createNode("TEST01", new ArrayList<>());
        
        // 执行
        NodeResult result = calculator.calculate(node, context);
        
        // 验证：应该返回空结果
        assertNotNull(result);
        assertEquals("TEST01", result.getStcd());
    }
    
    @Test
    void testCalculate_DirectInflow() {
        // 准备：河道直接入流
        TopoNode upstreamNode = createNode("UPSTREAM", new ArrayList<>());
        List<TimeSeriesData> upstreamData = createTimeSeriesData(3);
        upstreamData.get(0).setQSim(100.0);
        upstreamData.get(1).setQSim(150.0);
        upstreamData.get(2).setQSim(200.0);
        upstreamData.get(0).setQDsc(10.0);
        upstreamData.get(1).setQDsc(15.0);
        upstreamData.get(2).setQDsc(20.0);
        upstreamData.get(0).setQDvr(5.0);
        upstreamData.get(1).setQDvr(8.0);
        upstreamData.get(2).setQDvr(10.0);
        upstreamNode.setData(upstreamData);
        
        // 缓存上游结果
        context.cacheResult("UPSTREAM", new NodeResult("UPSTREAM", upstreamData));
        
        // 创建下游节点
        InputComponent input = new InputComponent();
        input.setArtp("D"); // 河道直接入流
        input.setArcd("UPSTREAM");
        input.setData(createTimeSeriesData(3));
        input.setMdcdList(Collections.singletonList(new ModelInfo()));
        
        TopoNode node = createNode("DOWNSTREAM", Collections.singletonList(input));
        node.setData(createTimeSeriesData(3));
        
        // 执行
        NodeResult result = calculator.calculate(node, context);
        
        // 验证：QSim = 上游QSim + QDsc - QDvr
        assertNotNull(result);
        // 注意：由于时间判断，只有在 startTime 之后的数据才会计算
    }
    
    @Test
    void testCalculate_ReservoirInflow() {
        // 准备：水库出流
        TopoNode reservoirNode = createNode("RESERVOIR", new ArrayList<>());
        List<TimeSeriesData> reservoirData = createTimeSeriesData(3);
        reservoirData.get(0).setQOut(80.0);
        reservoirData.get(1).setQOut(90.0);
        reservoirData.get(2).setQOut(100.0);
        reservoirNode.setData(reservoirData);
        
        // 缓存水库结果
        context.cacheResult("RESERVOIR", new NodeResult("RESERVOIR", reservoirData));
        
        // 创建下游节点
        InputComponent input = new InputComponent();
        input.setArtp("RSVR"); // 水库入流
        input.setArcd("RESERVOIR");
        input.setData(createTimeSeriesData(3));
        input.setMdcdList(Collections.singletonList(new ModelInfo()));
        
        TopoNode node = createNode("DOWNSTREAM", Collections.singletonList(input));
        node.setData(createTimeSeriesData(3));
        
        // 执行
        NodeResult result = calculator.calculate(node, context);
        
        // 验证
        assertNotNull(result);
        assertEquals("DOWNSTREAM", result.getStcd());
    }
    
    @Test
    void testSumInflows() {
        // 准备：多个输入分量
        InputComponent input1 = new InputComponent();
        input1.setData(createTimeSeriesData(3));
        input1.getData().get(0).setQSim(50.0);
        input1.getData().get(1).setQSim(60.0);
        input1.getData().get(2).setQSim(70.0);
        
        InputComponent input2 = new InputComponent();
        input2.setData(createTimeSeriesData(3));
        input2.getData().get(0).setQSim(30.0);
        input2.getData().get(1).setQSim(40.0);
        input2.getData().get(2).setQSim(50.0);
        
        List<InputComponent> inputs = Arrays.asList(input1, input2);
        TopoNode node = createNode("TEST", inputs);
        node.setData(createTimeSeriesData(3));
        
        // 执行
        NodeResult result = calculator.calculate(node, context);
        
        // 验证：QSim 应该是两个输入的和
        assertNotNull(result);
        // 注意：实际的累加在 calculate 方法内部完成
    }
    
    @Test
    void testAddBaseFlow() {
        // 准备：有基流的节点
        TopoNode node = createNode("TEST", new ArrayList<>());
        node.setData(createTimeSeriesData(3));
        node.getData().get(0).setQSim(100.0);
        node.getData().get(1).setQSim(150.0);
        node.getData().get(2).setQSim(200.0);
        
        // 设置基流
        List<Double> baseFlow = Arrays.asList(10.0, 15.0, 20.0);
        node.setQBases(baseFlow);
        
        // 执行
        NodeResult result = calculator.calculate(node, context);
        
        // 验证：QSim 应该叠加基流
        assertNotNull(result);
        // 基流在 calculate 方法内部叠加
    }
    
    @Test
    void testCalculateIntakes() {
        // 准备：有取水口的节点
        IntakeInfo intake1 = new IntakeInfo();
        intake1.setData(createTimeSeriesData(3));
        intake1.getData().get(0).setQ(5.0);
        intake1.getData().get(1).setQ(8.0);
        intake1.getData().get(2).setQ(10.0);
        
        IntakeInfo intake2 = new IntakeInfo();
        intake2.setData(createTimeSeriesData(3));
        intake2.getData().get(0).setQ(3.0);
        intake2.getData().get(1).setQ(4.0);
        intake2.getData().get(2).setQ(5.0);
        
        InputComponent input = new InputComponent();
        input.setArtp("D");
        input.setArcd("UPSTREAM");
        input.setData(createTimeSeriesData(3));
        input.setIntakes(Arrays.asList(intake1, intake2));
        input.setMdcdList(Collections.singletonList(new ModelInfo()));
        
        // 缓存上游结果
        TopoNode upstream = createNode("UPSTREAM", new ArrayList<>());
        upstream.setData(createTimeSeriesData(3));
        upstream.getData().forEach(d -> d.setQSim(100.0));
        context.cacheResult("UPSTREAM", new NodeResult("UPSTREAM", upstream.getData()));
        
        TopoNode node = createNode("TEST", Collections.singletonList(input));
        node.setData(createTimeSeriesData(3));
        
        // 执行
        NodeResult result = calculator.calculate(node, context);
        
        // 验证：应该扣除取水量
        assertNotNull(result);
    }
    
    @Test
    void testHubDistribution() {
        // 准备：枢纽节点（需要有输入分量才会执行完整流程）
        GateInfo gate = new GateInfo();
        gate.setGttp("DIS"); // 分流闸门
        gate.setDtcd("DIST01");
        
        HubInfo hub = new HubInfo();
        hub.setGate(Collections.singletonList(gate));
        
        // 创建一个空的输入分量（避免提前返回）
        InputComponent dummyInput = new InputComponent();
        dummyInput.setArtp("D");
        dummyInput.setData(createTimeSeriesData(3));
        // 没有 mdcdList，所以不会执行模型计算
        
        TopoNode node = createNode("HUB", Collections.singletonList(dummyInput));
        node.setHub(hub);
        node.setData(createTimeSeriesData(3));
        node.getData().forEach(d -> d.setQSim(100.0));
        
        // 执行
        NodeResult result = calculator.calculate(node, context);
        
        // 验证：应该记录枢纽分流数据
        assertNotNull(result);
        assertTrue(context.hasHubDistribution("DIST01"));
    }
    
    // ========== 辅助方法 ==========
    
    private TopoNode createNode(String stcd, List<InputComponent> inputs) {
        TopoNode node = new TopoNode();
        node.setStcd(stcd);
        node.setStnm("测试站点-" + stcd);
        node.setInputs(inputs);
        return node;
    }
    
    private List<TimeSeriesData> createTimeSeriesData(int count) {
        List<TimeSeriesData> dataList = new ArrayList<>();
        long now = System.currentTimeMillis();
        for (int i = 0; i < count; i++) {
            TimeSeriesData data = new TimeSeriesData();
            data.setDt(new Date(now + i * 3600000)); // 每小时
            data.setQ(0.0);
            data.setQSim(0.0);
            data.setQDvr(0.0);
            data.setQDsc(0.0);
            dataList.add(data);
        }
        return dataList;
    }
}

